package com.dbms.execution.execution;

import com.dbms.catalog.Catalog;
import com.dbms.common.Schema;
import com.dbms.common.Tuple;
import com.dbms.common.RID;
import com.dbms.execution.plan.PlanNode;
import com.dbms.parser.ast.ArithmeticNode;
import com.dbms.storage.BufferPoolManager;
import com.dbms.storage.Page;
import com.dbms.storage.TablePage;
import com.dbms.transaction.LogManager;
import com.dbms.transaction.LogRecord;
import com.dbms.transaction.Transaction;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * UpdateExecutor (更新执行器) 负责处理元组的更新操作。
 * 它首先使用索引找到要更新的元组，然后根据提供的更新映射（updates）
 * 创建一个新的元组版本，并用新版本替换旧版本。
 * 【已修改】这个版本增加了对算术更新的支持，例如 `SET salary = salary + 1000`。
 */
public class UpdateExecutor implements Executor {
    private final BufferPoolManager bufferPoolManager;
    private final Catalog.TableMetadata tableInfo;
    private final Integer keyToUpdate; // 要更新的元组的主键
    /**
     * 一个映射，键是列名，值是要更新的内容。
     * 值可以是具体的值 (e.g., "NewName")，也可以是一个 ArithmeticNode (e.g., for "salary + 1000")。
     */
    private final Map<String, Object> updates;
    private boolean hasExecuted = false;
    private final Transaction transaction;
    private final LogManager logManager;

    public UpdateExecutor(BufferPoolManager bpm, Catalog.TableMetadata tableInfo, Integer key, Map<String, Object> updates, Transaction transaction, LogManager logManager) {
        this.bufferPoolManager = bpm;
        this.tableInfo = tableInfo;
        this.keyToUpdate = key;
        this.updates = updates;
        this.transaction = transaction;
        this.logManager = logManager;
    }

    @Override
    public void open() { hasExecuted = false; }

    @Override
    public Tuple next() {
        if (hasExecuted) return null;
        Page page = null;
        try {
            // 1. 使用索引查找要更新的元组的 RID
            RID rid = tableInfo.index.find(keyToUpdate);
            if (rid == null) {
                // 如果找不到记录，则无需执行任何操作
                hasExecuted = true;
                return null;
            }
            // 2. 获取包含该元组的数据页
            page = bufferPoolManager.fetchPage(rid.getPageId());
            page.wLock(); // 加写锁，因为我们要修改页面内容
            try {
                TablePage tablePage = new TablePage(page);
                byte[] oldTupleData = tablePage.getTuple(rid.getSlotNum());
                if (oldTupleData == null) return null; // 元组可能已被删除

                Tuple oldTuple = new Tuple(oldTupleData, tableInfo.schema);

                // 3. 如果在事务中，记录更新前的旧元组值（用于回滚）
                if (transaction != null && logManager != null) {
                    LogRecord logRecord = new LogRecord(transaction.getTransactionId(), tableInfo.tableName, rid, oldTuple);
                    logManager.appendLogRecord(logRecord);
                }

                // 4. 构建新元组的值列表
                List<Object> newValues = new ArrayList<>();
                for (int i = 0; i < tableInfo.schema.getColumnCount(); i++) {
                    String colName = tableInfo.schema.getColumn(i).getColumnName();
                    Object newValue = oldTuple.getValue(i); // 默认情况下，值保持不变

                    // 检查此列是否在需要更新的列表中
                    if (updates.containsKey(colName)) {
                        Object updateInstruction = updates.get(colName);

                        // --- 【修改开始】 ---
                        // 检查更新指令是算术节点还是直接的值
                        if (updateInstruction instanceof ArithmeticNode) {
                            // 处理 'SET col = col + val' 这样的算术更新
                            ArithmeticNode node = (ArithmeticNode) updateInstruction;
                            Object oldValue = oldTuple.getValue(i);

                            if (!(oldValue instanceof Integer) || !(node.getValue() instanceof Integer)) {
                                throw new UnsupportedOperationException("Arithmetic updates are only supported for INTEGER columns.");
                            }

                            Integer oldIntValue = (Integer) oldValue;
                            Integer operand = (Integer) node.getValue();

                            // 根据操作符执行计算
                            if ("+".equals(node.getOperator())) {
                                newValue = oldIntValue + operand;
                            } else if ("-".equals(node.getOperator())) {
                                newValue = oldIntValue - operand;
                            } else {
                                throw new UnsupportedOperationException("Unsupported operator in UPDATE statement: " + node.getOperator());
                            }
                        } else {
                            // 处理 'SET col = val' 这样的直接赋值更新
                            newValue = updateInstruction;
                        }
                        // --- 【修改结束】 ---
                    }

                    // 检查是否尝试修改主键 (假设主键是第一列)
                    if (i == 0 && !newValue.equals(oldTuple.getValue(i))) {
                        throw new RuntimeException("Primary key cannot be updated.");
                    }

                    newValues.add(newValue);
                }
                // 5. 创建新的元组对象
                Tuple newTuple = new Tuple(newValues, tableInfo.schema);

                // 6. 在页面上执行更新操作
                tablePage.updateTuple(rid.getSlotNum(), newTuple.getData());

            } finally {
                // 7. 解锁并解钉页面，标记为脏页 (dirty=true)
                page.wUnlock();
                bufferPoolManager.unpinPage(page.getPageId(), true);
            }
            hasExecuted = true;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return null; // DML 操作不返回元组
    }

    @Override
    public void close() {}
    @Override
    public Schema getSchema() { return null; }

    @Override
    public PlanNode getPlanNode() {
        return null;
    }
}